from lib2to3.pytree import type_repr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from joblib import dump
from rdkit import DataStructs
from rdkit.Chem import AllChem, MACCSkeys
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.metrics import precision_score, classification_report, recall_score, f1_score, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
import logging
import torch
import torch.nn as nn

from base_model import BaseModle

REACTION_CLASS_NO="reaction class no."
ML_CLASS_NO = "ML class no."
SUBSTRATE_SMILES_CANONICAL= "substrate_smiles_canonical"


## 分箱或分区
def select_ml_class_no(t_reaction_class_no):
    if t_reaction_class_no in range(0,8):
        return t_reaction_class_no
    elif t_reaction_class_no in range(8,11):
        return 5
    elif t_reaction_class_no in range(12, 29):
        return  t_reaction_class_no-3
    elif t_reaction_class_no in range(29, 30):
        return 24
    elif t_reaction_class_no in range(31, 37):
        return t_reaction_class_no-6
    elif t_reaction_class_no in range(37, 39):
        return t_reaction_class_no-7
    elif t_reaction_class_no in range(39, 46):
        return  32
    else:
        return t_reaction_class_no

class ClassificationProject:
    def __init__(self):
        # pandas warnings
        pd.options.mode.chained_assignment = None  # default='warn'
        logging.basicConfig(filename='./classfication.log', level=logging.DEBUG)


    def handle(self,df):
                ## 计算y
        data,self.y = self.cal_data_y(df) 


        ## 计算x
        x1 = self.cal_x(data,"MACCS")
        x2 = self.cal_x(data,"MORGAN")
        self.x_append = np.append(x1, x2, axis=1)

        logging.info(self.x_append .shape)

        ## dump训练
        y_test,y_pred = self.tran_data()


        ## 输出图片
        self.export_img(y_test,y_pred)


    ## 读取原始数据
    def read_data(self,filename="database_0714.xlsx"):
        df = pd.read_excel(filename)
        # 先分类简化
        for i in range(len(df)):
            df[ML_CLASS_NO][i]  =  select_ml_class_no(int(df[REACTION_CLASS_NO][i]))   
        return df



 
    ## 这里只涉及一个class的特征，所以应该抽象一个函数
    def cal_data_y(self,df):
        ## 1. 组装成data结构
        data = df[["PUBCHEM CID", "source_substrate_smiles", "substrate_smiles_canonical"]]
        data.drop_duplicates(inplace=True)
        data = data.reset_index(drop=True)
        data[ML_CLASS_NO] = None
        logging.debug(data)
        for i in range(len(data)):
            smi = data[SUBSTRATE_SMILES_CANONICAL][i]
            lst = df[df[SUBSTRATE_SMILES_CANONICAL] == smi][ML_CLASS_NO].to_list()
            data[ML_CLASS_NO][i] = lst

        
        # transform reaction class number into numpy array for model
        arr_length = df[ML_CLASS_NO].max()


        ## 组装成（NXM）数组，记录下相关性特征。0~33种类别。 --》这里没明白，感觉不需要这样大的矩阵
        arr_index = np.zeros((len(data), arr_length))
        for i in range(len(data)):
            for t_class in data[ML_CLASS_NO][i]:
                arr_index[i][t_class-1] = 1 ## -1防止越界
        return data,arr_index.astype(float)
    



    def cal_x(self,data,type):
        ## 重复率过高
        ls = []
        for i in range(len(data)):
            smi = data[SUBSTRATE_SMILES_CANONICAL][i]
            try:
                m = AllChem.MolFromSmiles(smi)
                if type == "MACCS":
                    features = MACCSkeys.GenMACCSKeys(m)
                elif type == "MORGAN":
                    features = AllChem.GetMorganFingerprintAsBitVect(m, 3, nBits=2048)
                array = np.zeros((0,), dtype=np.float64)
                DataStructs.ConvertToNumpyArray(features, array)
                ls.append(array)
            except (Exception,):
                logging.error(smi)
                logging.error(i)
        X = np.array(ls)
        X = X.astype(float)
        return X


    def tran_data(self):
        base_model = BaseModle(self.x_append, self.y)
        ## 这里抽一个函数出来
        base_model.torch()
        base_model.mlp()
        base_model.clf()

        base_model.export_img()

        

        


if __name__ == "__main__":
    pj = ClassificationProject()
    df = pj.read_data()
    pj.handle(df)